import sys
sys.path.append('/usr/lib/python2.6/site-packages')

from numpy import *
import time

d = int(sys.argv[1]) #CUDA
n = int(sys.argv[2]) #Number of bodies
k = int(sys.argv[3]) #Number of iterations

G = 1     #Gravitational constant
dT = 0.01 #Time increment

M  = empty((n,1), cuda=d)
ufunc_random(M,M)
MT = empty((1,n), cuda=d)
ufunc_random(MT,MT)
Px = empty((n,1), cuda=d)
ufunc_random(Px,Px)
Py = empty((n,1), cuda=d)
ufunc_random(Py,Py)
Pz = empty((n,1), cuda=d)
ufunc_random(Pz,Pz)
PxT= empty((1,n), cuda=d)
ufunc_random(PxT,PxT)
PyT= empty((1,n), cuda=d)
ufunc_random(PyT,PyT)
PzT= empty((1,n), cuda=d)
ufunc_random(PzT,PzT)
Vx = empty((n,1), cuda=d)
ufunc_random(Vx,Vx)
Vy = empty((n,1), cuda=d)
ufunc_random(Vy,Vy)
Vz = empty((n,1), cuda=d)
ufunc_random(Vz,Vz)

OnesCol = zeros((n,1), dtype=double, cuda=d)+1.0
OnesRow = zeros((1,n), dtype=double, cuda=d)+1.0
#Identity= array(diag([1]*n), dtype=double, cuda=d)
rt=0
stime = time.time()
for i in range(k):
    timing = False
    if timing:
        print '-+-+-+-+-+-+-+-+-+-+-+-+-'
        t0=time.time()
    #distance between all pairs of objects
    Fx = dot(OnesCol, PxT) - dot(Px, OnesRow)
    Fy = dot(OnesCol, PyT) - dot(Py, OnesRow)
    Fz = dot(OnesCol, PzT) - dot(Pz, OnesRow)
    if timing:
        t1=time.time()
        print 'a time:', t1-t0
        t0=time.time()
    Dsq = Fx * Fx + Fy * Fy + Fx * Fz #+ Identity
    D = sqrt(Dsq)
    if timing:
        t1=time.time()
        print 'b time:', t1-t0
        t0=time.time()
    #mutual forces between all pairs of objects
    F = G * dot(M, MT) / Dsq
    if timing:
        t1=time.time()
        print 'c time:', t1-t0
        t0=time.time()
    #F = F - diag(diag(F))#set 'self attraction' to 0
    Fx = (Fx / D) * F
    Fy = (Fy / D) * F
    Fz = (Fz / D) * F
    if timing:
        t1=time.time()
        print 'd time:', t1-t0
    addt0=time.time()
    #net force on each body
    Fnet_x = add.reduce(Fx,1)
    Fnet_y = add.reduce(Fy,1)
    Fnet_z = add.reduce(Fz,1)
    addt1=time.time()
    rt += addt1-addt0
    if timing:
        t1=time.time()
        print 'e time:', t1-t0
        t0=time.time()
    #print shape(Fnet_x)
    Fnet_x = Fnet_x[:,newaxis]
    Fnet_y = Fnet_y[:,newaxis]
    Fnet_z = Fnet_z[:,newaxis]
    #print shape(Fnet_x)
    if timing:
        t1=time.time()
        print 'f time:', t1-t0
        t0=time.time()
    Fnet_x_tmp = Fnet_x*dT
    Fnet_y_tmp = Fnet_y*dT
    Fnet_z_tmp = Fnet_z*dT 
    if timing:
        t1=time.time()
        print 'g time:', t1-t0
        t0=time.time()
    #change in velocity:
    Vx += Fnet_x_tmp / M
    Vy += Fnet_y_tmp / M
    Vz += Fnet_z_tmp / M
    if timing:
        t1=time.time()
        print 'h time:', t1-t0
        t0=time.time()
    #change in position
    Px += Vx * dT
    Py += Vy * dT
    Pz += Vz * dT
    if timing:
        t1=time.time()
        print 'i time:', t1-t0

print 'nbody with #bodies: ', n,', iter: ', i+1, 'in sec: ', time.time() - stime, 'reduce took:',rt/k

